# uamcf/utils/io.py
import os
import json
import yaml
from typing import Dict, Any, Optional, Union

from .logger import get_logger

logger = get_logger("IO")


def load_config(config_path: str) -> Dict[str, Any]:
    """
    加载配置文件

    Args:
        config_path: 配置文件路径 (支持JSON和YAML)

    Returns:
        配置字典
    """
    if not os.path.exists(config_path):
        logger.error(f"Config file not found: {config_path}")
        return {}

    try:
        file_ext = os.path.splitext(config_path)[1].lower()

        if file_ext in ('.json'):
            with open(config_path, 'r') as f:
                config = json.load(f)
        elif file_ext in ('.yaml', '.yml'):
            try:
                import yaml
                with open(config_path, 'r') as f:
                    config = yaml.safe_load(f)
            except ImportError:
                logger.error("PyYAML is not installed. Please install with 'pip install pyyaml'")
                return {}
        else:
            logger.error(f"Unsupported config file format: {file_ext}")
            return {}

        logger.info(f"Config loaded from {config_path}")
        return config

    except Exception as e:
        logger.error(f"Failed to load config: {str(e)}")
        return {}


def save_config(config: Dict[str, Any], output_path: str) -> bool:
    """
    保存配置到文件

    Args:
        config: 配置字典
        output_path: 输出文件路径 (支持JSON和YAML)

    Returns:
        是否成功保存
    """
    try:
        # 确保目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        file_ext = os.path.splitext(output_path)[1].lower()

        if file_ext in ('.json'):
            with open(output_path, 'w') as f:
                json.dump(config, f, indent=2)
        elif file_ext in ('.yaml', '.yml'):
            try:
                import yaml
                with open(output_path, 'w') as f:
                    yaml.dump(config, f, default_flow_style=False)
            except ImportError:
                logger.error("PyYAML is not installed. Please install with 'pip install pyyaml'")
                return False
        else:
            logger.error(f"Unsupported config file format: {file_ext}")
            return False

        logger.info(f"Config saved to {output_path}")
        return True

    except Exception as e:
        logger.error(f"Failed to save config: {str(e)}")
        return False


def load_dataset(dataset_path: str, framework: str = "pytorch") -> Optional[Any]:
    """
    加载数据集

    Args:
        dataset_path: 数据集路径
        framework: 使用的框架 ('pytorch', 'tensorflow')

    Returns:
        加载的数据集对象，若加载失败则返回None
    """
    try:
        if framework.lower() == "pytorch":
            try:
                import torch
                from torch.utils.data import Dataset, DataLoader

                # 尝试加载torch格式数据集
                if os.path.isfile(dataset_path):
                    return torch.load(dataset_path)

                # 尝试加载目录形式的数据集
                elif os.path.isdir(dataset_path):
                    logger.warning("Directory-based datasets require custom loading logic")
                    return None

                else:
                    logger.error(f"Dataset path not found: {dataset_path}")
                    return None

            except ImportError:
                logger.error("PyTorch is not installed. Please install with 'pip install torch'")
                return None

        elif framework.lower() == "tensorflow":
            try:
                import tensorflow as tf

                # 尝试加载SavedModel格式
                if os.path.isdir(dataset_path):
                    return tf.data.experimental.load(dataset_path)

                # 尝试加载TFRecord格式
                elif dataset_path.endswith('.tfrecord'):
                    return tf.data.TFRecordDataset(dataset_path)

                else:
                    logger.error(f"Unsupported dataset format for TensorFlow: {dataset_path}")
                    return None

            except ImportError:
                logger.error("TensorFlow is not installed. Please install with 'pip install tensorflow'")
                return None

        else:
            logger.error(f"Unsupported framework: {framework}")
            return None

    except Exception as e:
        logger.error(f"Failed to load dataset: {str(e)}")
        return None


def export_results(results: Dict[str, Any], output_path: str, format: str = "json") -> bool:
    """
    导出压缩结果

    Args:
        results: 结果字典
        output_path: 输出文件路径
        format: 输出格式 ('json', 'yaml', 'csv')

    Returns:
        是否成功导出
    """
    try:
        # 确保目录存在
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        if format.lower() == "json":
            with open(output_path, 'w') as f:
                json.dump(results, f, indent=2)

        elif format.lower() == "yaml":
            try:
                import yaml
                with open(output_path, 'w') as f:
                    yaml.dump(results, f, default_flow_style=False)
            except ImportError:
                logger.error("PyYAML is not installed. Please install with 'pip install pyyaml'")
                return False

        elif format.lower() == "csv":
            try:
                import csv
                import pandas as pd

                # 将嵌套字典扁平化
                flat_dict = {}

                def flatten_dict(d, parent_key=''):
                    for k, v in d.items():
                        key = f"{parent_key}.{k}" if parent_key else k
                        if isinstance(v, dict):
                            flatten_dict(v, key)
                        else:
                            flat_dict[key] = v

                flatten_dict(results)

                # 导出为CSV
                df = pd.DataFrame([flat_dict])
                df.to_csv(output_path, index=False)

            except ImportError:
                logger.error("Pandas is not installed. Please install with 'pip install pandas'")
                return False

        else:
            logger.error(f"Unsupported export format: {format}")
            return False

        logger.info(f"Results exported to {output_path}")
        return True

    except Exception as e:
        logger.error(f"Failed to export results: {str(e)}")
        return False